#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <math.h>
#include "bernoulli.h"
#include "random.h"

void Bernoulli_Init_Params(Bernoulli_t * dist, int nparams, ...)
{
  if (1 == nparams)
  {
    va_list ap;
  
    va_start(ap, nparams);

    dist->mean = va_arg(ap, double);
    dist->fixed_mean = 1;

    va_end(ap);

    if ((dist->mean < 0) || (dist->mean > 1))
    {
      fprintf(stderr, "error: Bernoulli probability must be in [0,1]\n");
      exit(1);
    }
  }
  else if (0 == nparams)
  {
    dist->fixed_mean = 0;
  }
  else
  {
    fprintf(stderr, "error: Bernoulli takes at most 1 parameter\n");
    exit(1);
  }
}


int Bernoulli_Sample(const Bernoulli_t * dist, int nargs, ...)
{
  double mean;
  
  if (dist->fixed_mean)
    mean = dist->mean;

  else if (1 == nargs)
  {
    va_list ap;
  
    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);

    if ((mean < 0) || (mean > 1))
    {
      fprintf(stderr, "error: Bernoulli probability must be in [0,1]\n");
      exit(1);
    }
  }
  else
  {
    fprintf(stderr, "error: Bernoulli requires 1 argument\n");
    exit(1);
  }

  if (RANDOM_DOUBLE < mean)
    return 0;                                /* true */
  else
    return 1;                                /* false */
}


double Bernoulli_LogProb(const Bernoulli_t * dist, int val, int nargs, ...)
{
  double mean;
  
  if (dist->fixed_mean)
    mean = dist->mean;

  else if (1 == nargs)
  {
    va_list ap;
  
    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);

    if ((mean < 0) || (mean > 1))
    {
      fprintf(stderr, "error: Bernoulli probability must be in [0,1]\n");
      exit(1);
    }
  }
  else
  {
    fprintf(stderr, "error: Bernoulli requires 1 argument\n");
    exit(1);
  }

  /* true */
  if (0 == val)
    return log(mean);

  /* false */
  else if (1 == val)
    return log(1 - mean);

  else
    return -HUGE_VAL;
}

double Bernoulli_Prob(const Bernoulli_t * dist, int val, int nargs, ...)
{
  double mean;
  
  if (dist->fixed_mean)
    mean = dist->mean;

  else if (1 == nargs)
  {
    va_list ap;
  
    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);
  }
  else
  {
    fprintf(stderr, "error: Bernoulli requires 1 argument\n");
    exit(1);
  }

  /* true */
  if (0 == val)
    return mean;

  /* false */
  else if (1 == val)
    return 1 - mean;

  else
    return 0.0;
}
